#pragma once

#include <cstddef>

#include "../utils/misc.h"

namespace mdtx::online::impl
{
    template <typename T>
    class var
    {
    public:
        ptrdiff_t n;
        T m, s2;
        kahan<T> acc_m, acc_s2;

        void insert_impl(T x)
        {
            T d = x - m;
            n += 1;
            acc_m(m, d / n);
            acc_s2(s2, (x - m) * d);
        }

        void remove_impl(T x)
        {
            T d = m - x;
            n -= 1;
            acc_m(m, d / n);
            acc_s2(s2, (x - m) * d);
        }

        void reset_impl()
        {
            n = 0;
            m = s2 = 0;
            acc_m.reset();
            acc_s2.reset();
        }
    };

    template <typename T>
    class var_vw
    {
    public:
        T wsum, w2sum;
        T m, s2;
        kahan<T> acc_w, acc_w2, acc_m, acc_s2;

        void insert_impl(T x, T w)
        {
            T d = x - m;
            acc_w(wsum, w);
            acc_w2(w2sum, w * w);
            acc_m(m, (w / wsum) * d);
            acc_s2(s2, w * (x - m) * d);
        }

        void remove_impl(T x, T w)
        {
            T d = m - x;
            acc_w(wsum, -w);
            acc_w2(w2sum, -w * w);
            acc_m(m, (w / wsum) * d);
            acc_s2(s2, w * (x - m) * d);
        }

        void reset_impl()
        {
            wsum = w2sum = m = s2 = 0;
            acc_w.reset();
            acc_w2.reset();
            acc_m.reset();
            acc_s2.reset();
        }
    };

    template <typename T>
    class var_ew
    {
    public:
        T a, b;
        T m{}, s2{};
        kahan<T> acc_m{};
        bool init{};

        explicit var_ew(T alpha)
            : a(alpha), b(1.0 - a)
        {
        }

        void insert_impl(T x)
        {
            if (init)
                [[likely]]
            {
                T d = x - m;
                T i = a * d;
                acc_m(m, i);
                s2 = b * (s2 + i * d);
            }
            else
                [[unlikely]]
            {
                m = x;
                init = true;
            }
        }

        void reset_impl()
        {
            init = false;
            m = s2 = 0;
            acc_m.reset();
        }
    };

    template <typename T>
    class var_evw
    {
    public:
        T m, s2;
        kahan<T> acc_m;
        bool init;

        void insert_impl(T x, T w)
        {
            if (init)
                [[likely]]
            {
                T d = x - m;
                T i = w * d;
                acc_m(m, i);
                s2 = (1 - w) * (s2 + i * d);
            }
            else
                [[unlikely]]
            {
                m = x;
                init = true;
            }
        }

        void reset_impl()
        {
            init = false;
            m = s2 = 0;
            acc_m.reset();
        }
    };

    template <typename T>
    class covxy
    {
    public:
        ptrdiff_t n;
        T mx, my, sxy;
        kahan<T> acc_mx, acc_my, acc_sxy;

        void insert_impl(T x, T y)
        {
            n += 1;
            T dx = x - mx;
            T dy = y - my;
            acc_mx(mx, dx / n);
            acc_my(my, dy / n);
            acc_sxy(sxy, (x - mx) * dy);
        }

        void remove_impl(T x, T y)
        {
            n -= 1;
            T dx = mx - x;
            T dy = my - y;
            acc_mx(mx, dx / n);
            acc_my(my, dy / n);
            acc_sxy(sxy, (x - mx) * dy);
        }

        void reset_impl()
        {
            n = 0;
            mx = my = sxy = 0;
            acc_mx.reset();
            acc_my.reset();
            acc_sxy.reset();
        }
    };

    template <typename T>
    class covxy_vw
    {
    public:
        T wsum, w2sum;
        T mx, my, sxy;
        kahan<T> acc_w, acc_w2, acc_mx, acc_my, acc_sxy;

        void insert_impl(T x, T y, T w)
        {
            acc_w(wsum, w);
            acc_w2(w2sum, w * w);
            T dx = x - mx;
            T dy = y - my;
            acc_mx(mx, (w / wsum) * dx);
            acc_my(my, (w / wsum) * dy);
            acc_sxy(sxy, w * (x - mx) * dy);
        }

        void remove_impl(T x, T y, T w)
        {
            acc_w(wsum, -w);
            acc_w2(w2sum, -w * w);
            T dx = mx - x;
            T dy = my - y;
            acc_mx(mx, (w / wsum) * dx);
            acc_my(my, (w / wsum) * dy);
            acc_sxy(sxy, w * (x - mx) * dy);
        }

        void reset_impl()
        {
            wsum = w2sum = mx = my = sxy = 0;
            acc_w.reset();
            acc_w2.reset();
            acc_mx.reset();
            acc_my.reset();
            acc_sxy.reset();
        }
    };

    template <typename T>
    class covxy_ew
    {
    public:
        bool init{};
        T alpha, beta;
        T mx{}, my{}, sxy{};
        kahan<T> acc_mx{}, acc_my{};

        explicit covxy_ew(T alpha)
            : alpha(alpha), beta(1.0 - alpha)
        {
        }

        void insert_impl(T x, T y)
        {
            if (init)
                [[likely]]
            {
                T dx = x - mx;
                T dy = y - my;
                T ix = alpha * dx;
                T iy = alpha * dy;
                acc_mx(mx, ix);
                acc_my(my, iy);
                sxy = beta * (sxy + dx * iy);
            }
            else
                [[unlikely]]
            {
                mx = x;
                my = y;
                init = true;
            }
        }

        void reset_impl()
        {
            init = false;
            mx = my = sxy = 0;
            acc_mx.reset();
            acc_my.reset();
        }
    };

    template <typename T>
    class corrxy
    {
    public:
        ptrdiff_t n;
        T mx, my, sx, sy, sxsy, sxy;
        kahan<T> acc_mx, acc_my, acc_sx, acc_sy, acc_sxy;

        void insert_impl(T x, T y)
        {
            n += 1;
            T dx = x - mx;
            T dy = y - my;
            acc_mx(mx, dx / n);
            acc_my(my, dy / n);
            acc_sx(sx, (x - mx) * dx);
            acc_sy(sy, (y - my) * dy);
            acc_sxy(sxy, (x - mx) * dy);
            sxsy = sx * sy;
        }

        void remove_impl(T x, T y)
        {
            n -= 1;
            T dx = mx - x;
            T dy = my - y;
            acc_mx(mx, dx / n);
            acc_my(my, dy / n);
            acc_sx(sx, (x - mx) * dx);
            acc_sy(sy, (y - my) * dy);
            acc_sxy(sxy, (x - mx) * dy);
            sxsy = sx * sy;
        }

        void reset_impl()
        {
            n = 0;
            mx = my = sx = sy = sxsy = sxy = 0;
            acc_mx.reset();
            acc_my.reset();
            acc_sx.reset();
            acc_sy.reset();
            acc_sxy.reset();
        }
    };

    template <typename T>
    class corrxy_vw
    {
    public:
        T wsum, w2sum;
        T mx, my, sx, sy, sxsy, sxy;
        kahan<T> acc_w, acc_w2, acc_mx, acc_my, acc_sx, acc_sy, acc_sxy;

        void insert_impl(T x, T y, T w)
        {
            acc_w(wsum, w);
            acc_w2(w2sum, w * w);
            T dx = x - mx;
            T dy = y - my;
            acc_mx(mx, (w / wsum) * dx);
            acc_my(my, (w / wsum) * dy);
            acc_sx(sx, w * (x - mx) * dx);
            acc_sy(sy, w * (y - my) * dy);
            acc_sxy(sxy, w * (x - mx) * dy);
            sxsy = sx * sy;
        }

        void remove_impl(T x, T y, T w)
        {
            acc_w(wsum, -w);
            acc_w2(w2sum, -w * w);
            T dx = mx - x;
            T dy = my - y;
            acc_mx(mx, (w / wsum) * dx);
            acc_my(my, (w / wsum) * dy);
            acc_sx(sx, w * (x - mx) * dx);
            acc_sy(sy, w * (y - my) * dy);
            acc_sxy(sxy, w * (x - mx) * dy);
            sxsy = sx * sy;
        }

        void reset_impl()
        {
            wsum = w2sum = mx = my = sx = sy = sxsy = sxy = 0;
            acc_w.reset();
            acc_w2.reset();
            acc_mx.reset();
            acc_my.reset();
            acc_sx.reset();
            acc_sy.reset();
            acc_sxy.reset();
        }
    };

    template <typename T>
    class corrxy_ew
    {
    public:
        bool init{};
        T alpha, beta;
        T mx{}, my{}, sx{}, sy{}, sxsy{}, sxy{};
        kahan<T> acc_mx{}, acc_my{};

        explicit corrxy_ew(T alpha)
            : alpha(alpha), beta(1 - alpha)
        {
        }

        void insert_impl(T x, T y)
        {
            if (init)
                [[likely]]
            {
                T dx = x - mx;
                T dy = y - my;
                T ix = alpha * dx;
                T iy = alpha * dy;
                acc_mx(mx, ix);
                acc_my(my, iy);
                sx = beta * (sx + dx * ix);
                sy = beta * (sy + dy * iy);
                sxy = beta * (sxy + dx * iy);
                sxsy = sx * sy;
            }
            else
                [[unlikely]]
            {
                mx = x;
                my = y;
                init = true;
            }
        }

        void reset_impl()
        {
            init = false;
            mx = my = sx = sy = sxsy = sxy = 0;
            acc_mx.reset();
            acc_my.reset();
        }
    };

    template <typename T>
    class beta
    {
    public:
        ptrdiff_t n;
        T mx, my, sx, sxy;
        kahan<T> acc_mx, acc_my, acc_sx, acc_sxy;

        void insert_impl(T x, T y)
        {
            n += 1;
            T dx = x - mx;
            T dy = y - my;
            acc_mx(mx, dx / n);
            acc_my(my, dy / n);
            acc_sx(sx, (x - mx) * dx);
            acc_sxy(sxy, (x - mx) * dy);
        }

        void remove_impl(T x, T y)
        {
            n -= 1;
            T dx = mx - x;
            T dy = my - y;
            acc_mx(mx, dx / n);
            acc_my(my, dy / n);
            acc_sx(sx, (x - mx) * dx);
            acc_sxy(sxy, (x - mx) * dy);
        }

        void reset_impl()
        {
            n = 0;
            mx = my = sx = sxy = 0;
            acc_mx.reset();
            acc_my.reset();
            acc_sx.reset();
            acc_sxy.reset();
        }
    };

    template <typename T>
    inline T rel_weight(T wsum, T w2sum)
    {
        return wsum - w2sum / wsum;
    }
}
